import asyncio
import websockets
import subprocess
import threading
import json
import base64
import cv2
import numpy as np
import io
from PIL import Image
from backend.services.device_manager.device_manager import DeviceManager
import logging
import os

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger('websocket_service')

class WebSocketService:
    _instance = None
    _lock = threading.Lock()
    
    def __new__(cls):
        with cls._lock:
            if cls._instance is None:
                cls._instance = super(WebSocketService, cls).__new__(cls)
                cls._instance.device_manager = DeviceManager()
                cls._instance.active_connections = {}
                cls._instance.screen_streams = {}
                cls._instance.command_handlers = {
                    'tap': cls._instance._handle_tap,
                    'swipe': cls._instance._handle_swipe,
                    'text': cls._instance._handle_text,
                    'key': cls._instance._handle_key,
                    'start_screen_stream': cls._instance._handle_start_screen_stream,
                    'stop_screen_stream': cls._instance._handle_stop_screen_stream
                }
        return cls._instance
    
    async def start(self, host='0.0.0.0', port=9002):
        """启动WebSocket服务"""
        try:
            self.server = await websockets.serve(self._handle_connection, host, port)
            logger.info(f"WebSocket服务已启动，监听 {host}:{port}")
            await self.server.wait_closed()
        except Exception as e:
            logger.error(f"启动WebSocket服务时发生错误: {str(e)}")
    
    async def stop(self):
        """停止WebSocket服务"""
        try:
            if hasattr(self, 'server'):
                self.server.close()
                await self.server.wait_closed()
                logger.info("WebSocket服务已停止")
        except Exception as e:
            logger.error(f"停止WebSocket服务时发生错误: {str(e)}")
    
    async def _handle_connection(self, websocket, path):
        """处理WebSocket连接"""
        # 生成唯一的连接ID
        connection_id = self._generate_connection_id()
        self.active_connections[connection_id] = websocket
        logger.info(f"新的WebSocket连接已建立: {connection_id}")
        
        try:
            # 接收客户端消息
            async for message in websocket:
                await self._process_message(connection_id, message)
        except websockets.ConnectionClosed as e:
            logger.info(f"WebSocket连接已关闭: {connection_id}, 原因: {str(e)}")
        except Exception as e:
            logger.error(f"处理WebSocket消息时发生错误: {str(e)}")
        finally:
            # 清理资源
            if connection_id in self.active_connections:
                del self.active_connections[connection_id]
            
            # 停止相关的屏幕流
            for device_id, stream_info in list(self.screen_streams.items()):
                if connection_id in stream_info['connections']:
                    stream_info['connections'].remove(connection_id)
                    if not stream_info['connections']:
                        # 没有客户端连接，停止屏幕流
                        self._stop_screen_stream(device_id)
    
    async def _process_message(self, connection_id, message):
        """处理客户端消息"""
        try:
            # 解析JSON消息
            data = json.loads(message)
            
            # 获取消息类型和设备ID
            message_type = data.get('type')
            device_id = data.get('device_id')
            
            # 根据消息类型调用相应的处理函数
            if message_type in self.command_handlers:
                await self.command_handlers[message_type](connection_id, device_id, data)
            else:
                logger.warning(f"未知的消息类型: {message_type}")
        except json.JSONDecodeError:
            logger.error("接收到无效的JSON消息")
        except Exception as e:
            logger.error(f"处理消息时发生错误: {str(e)}")
    
    async def _handle_tap(self, connection_id, device_id, data):
        """处理点击事件"""
        try:
            x = data.get('x')
            y = data.get('y')
            
            # 验证坐标
            if x is None or y is None:
                logger.warning("点击事件缺少坐标信息")
                return
            
            # 执行点击命令
            command = f'input tap {x} {y}'
            result = self.device_manager.execute_command(device_id, command)
            
            # 发送响应
            await self._send_response(connection_id, {
                'type': 'tap_response',
                'success': result['success'],
                'device_id': device_id
            })
        except Exception as e:
            logger.error(f"处理点击事件时发生错误: {str(e)}")
    
    async def _handle_swipe(self, connection_id, device_id, data):
        """处理滑动事件"""
        try:
            x1 = data.get('x1')
            y1 = data.get('y1')
            x2 = data.get('x2')
            y2 = data.get('y2')
            duration = data.get('duration', 300)  # 默认300毫秒
            
            # 验证坐标
            if None in [x1, y1, x2, y2]:
                logger.warning("滑动事件缺少坐标信息")
                return
            
            # 执行滑动命令
            command = f'input swipe {x1} {y1} {x2} {y2} {duration}'
            result = self.device_manager.execute_command(device_id, command)
            
            # 发送响应
            await self._send_response(connection_id, {
                'type': 'swipe_response',
                'success': result['success'],
                'device_id': device_id
            })
        except Exception as e:
            logger.error(f"处理滑动事件时发生错误: {str(e)}")
    
    async def _handle_text(self, connection_id, device_id, data):
        """处理文本输入事件"""
        try:
            text = data.get('text')
            
            # 验证文本
            if not text:
                logger.warning("文本输入事件缺少文本信息")
                return
            
            # 执行文本输入命令
            command = f'input text "{text}"'  # 使用引号包裹文本，避免空格问题
            result = self.device_manager.execute_command(device_id, command)
            
            # 发送响应
            await self._send_response(connection_id, {
                'type': 'text_response',
                'success': result['success'],
                'device_id': device_id
            })
        except Exception as e:
            logger.error(f"处理文本输入事件时发生错误: {str(e)}")
    
    async def _handle_key(self, connection_id, device_id, data):
        """处理按键事件"""
        try:
            key_code = data.get('key_code')
            
            # 验证按键码
            if key_code is None:
                logger.warning("按键事件缺少按键码信息")
                return
            
            # 执行按键命令
            command = f'input keyevent {key_code}'
            result = self.device_manager.execute_command(device_id, command)
            
            # 发送响应
            await self._send_response(connection_id, {
                'type': 'key_response',
                'success': result['success'],
                'device_id': device_id
            })
        except Exception as e:
            logger.error(f"处理按键事件时发生错误: {str(e)}")
    
    async def _handle_start_screen_stream(self, connection_id, device_id, data):
        """处理开始屏幕流请求"""
        try:
            # 检查设备是否存在
            device = self.device_manager.get_device_by_id(device_id)
            if not device:
                logger.warning(f"设备不存在: {device_id}")
                await self._send_response(connection_id, {
                    'type': 'stream_response',
                    'success': False,
                    'message': '设备不存在',
                    'device_id': device_id
                })
                return
            
            # 检查设备是否在线
            if device.status != '在线':
                logger.warning(f"设备不在线: {device_id}")
                await self._send_response(connection_id, {
                    'type': 'stream_response',
                    'success': False,
                    'message': '设备不在线',
                    'device_id': device_id
                })
                return
            
            # 如果屏幕流不存在，启动新的屏幕流
            if device_id not in self.screen_streams:
                self._start_screen_stream(device_id)
            
            # 将连接添加到屏幕流的连接列表中
            self.screen_streams[device_id]['connections'].add(connection_id)
            
            logger.info(f"连接 {connection_id} 已订阅设备 {device_id} 的屏幕流")
            
            # 发送响应
            await self._send_response(connection_id, {
                'type': 'stream_response',
                'success': True,
                'message': '屏幕流已启动',
                'device_id': device_id
            })
        except Exception as e:
            logger.error(f"处理开始屏幕流请求时发生错误: {str(e)}")
            await self._send_response(connection_id, {
                'type': 'stream_response',
                'success': False,
                'message': str(e),
                'device_id': device_id
            })
    
    async def _handle_stop_screen_stream(self, connection_id, device_id, data):
        """处理停止屏幕流请求"""
        try:
            # 检查屏幕流是否存在
            if device_id not in self.screen_streams:
                logger.warning(f"屏幕流不存在: {device_id}")
                return
            
            # 从连接列表中移除
            if connection_id in self.screen_streams[device_id]['connections']:
                self.screen_streams[device_id]['connections'].remove(connection_id)
                logger.info(f"连接 {connection_id} 已取消订阅设备 {device_id} 的屏幕流")
            
            # 如果没有连接，停止屏幕流
            if not self.screen_streams[device_id]['connections']:
                self._stop_screen_stream(device_id)
        except Exception as e:
            logger.error(f"处理停止屏幕流请求时发生错误: {str(e)}")
    
    def _start_screen_stream(self, device_id):
        """启动屏幕流"""
        try:
            # 创建线程安全的连接集合
            connections = set()
            
            # 创建停止事件
            stop_event = threading.Event()
            
            # 启动屏幕流线程
            stream_thread = threading.Thread(target=self._stream_screen, args=(device_id, connections, stop_event))
            stream_thread.daemon = True
            stream_thread.start()
            
            # 保存屏幕流信息
            self.screen_streams[device_id] = {
                'connections': connections,
                'stop_event': stop_event,
                'thread': stream_thread
            }
            
            logger.info(f"设备 {device_id} 的屏幕流已启动")
        except Exception as e:
            logger.error(f"启动屏幕流时发生错误: {str(e)}")
    
    def _stop_screen_stream(self, device_id):
        """停止屏幕流"""
        try:
            # 检查屏幕流是否存在
            if device_id not in self.screen_streams:
                return
            
            # 获取屏幕流信息
            stream_info = self.screen_streams[device_id]
            
            # 设置停止事件
            stream_info['stop_event'].set()
            
            # 等待线程结束
            stream_info['thread'].join(timeout=5)
            
            # 移除屏幕流信息
            del self.screen_streams[device_id]
            
            logger.info(f"设备 {device_id} 的屏幕流已停止")
        except Exception as e:
            logger.error(f"停止屏幕流时发生错误: {str(e)}")
    
    def _stream_screen(self, device_id, connections, stop_event):
        """屏幕流线程函数"""
        try:
            # 启动scrcpy命令行工具，捕获屏幕输出
            scrcpy_cmd = [
                'scrcpy',
                '-s', device_id,
                '--no-display',
                '--bit-rate', '2M',
                '--max-size', '720',
                '--record', '-'
            ]
            
            # 启动进程
            process = subprocess.Popen(
                scrcpy_cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                bufsize=0
            )
            
            # 创建OpenCV视频捕获对象
            cap = cv2.VideoCapture(process.stdout.fileno())
            
            while not stop_event.is_set():
                # 读取一帧
                ret, frame = cap.read()
                if not ret:
                    break
                
                # 压缩图像
                _, buffer = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 80])
                
                # 将图像转换为base64编码
                jpg_as_text = base64.b64encode(buffer).decode('utf-8')
                
                # 创建消息
                message = json.dumps({
                    'type': 'screen_frame',
                    'device_id': device_id,
                    'frame': jpg_as_text
                })
                
                # 向所有连接的客户端发送帧
                # 注意：这里需要在异步事件循环中执行
                asyncio.run_coroutine_threadsafe(self._broadcast_frame(message, connections), asyncio.get_event_loop())
                
                # 限制帧率
                if stop_event.wait(timeout=0.033):  # 大约30fps
                    break
        except Exception as e:
            logger.error(f"屏幕流线程发生错误: {str(e)}")
        finally:
            # 清理资源
            try:
                if 'process' in locals() and process.poll() is None:
                    process.terminate()
                    process.wait(timeout=2)
            except:
                pass
            
            try:
                if 'cap' in locals():
                    cap.release()
            except:
                pass
    
    async def _broadcast_frame(self, message, connections):
        """向所有连接的客户端广播帧"""
        # 创建要移除的连接列表
        to_remove = []
        
        # 向每个连接发送消息
        for connection_id in connections:
            # 检查连接是否存在
            if connection_id not in self.active_connections:
                to_remove.append(connection_id)
                continue
            
            websocket = self.active_connections[connection_id]
            try:
                await websocket.send(message)
            except Exception as e:
                logger.warning(f"发送帧到连接 {connection_id} 时发生错误: {str(e)}")
                to_remove.append(connection_id)
        
        # 移除无效连接
        for connection_id in to_remove:
            connections.discard(connection_id)
    
    async def _send_response(self, connection_id, response):
        """发送响应消息"""
        # 检查连接是否存在
        if connection_id not in self.active_connections:
            return
        
        websocket = self.active_connections[connection_id]
        try:
            await websocket.send(json.dumps(response))
        except Exception as e:
            logger.warning(f"发送响应到连接 {connection_id} 时发生错误: {str(e)}")
    
    def _generate_connection_id(self):
        """生成唯一的连接ID"""
        # 使用UUID生成唯一ID
        return str(uuid.uuid4())
    
    def get_active_streams(self):
        """获取所有活动的屏幕流"""
        return {
            device_id: {
                'connection_count': len(stream_info['connections'])
            }
            for device_id, stream_info in self.screen_streams.items()
        }
    
    def execute_multi_device_command(self, device_ids, command_type, command_data):
        """在多台设备上执行命令（一机多控）"""
        results = {}
        
        for device_id in device_ids:
            try:
                # 根据命令类型执行不同的命令
                if command_type == 'tap':
                    command = f'input tap {command_data.get("x")} {command_data.get("y")}'
                elif command_type == 'swipe':
                    command = f'input swipe {command_data.get("x1")} {command_data.get("y1")} {command_data.get("x2")} {command_data.get("y2")} {command_data.get("duration", 300)}'
                elif command_type == 'text':
                    command = f'input text "{command_data.get("text")}"'  # 使用引号包裹文本，避免空格问题
                elif command_type == 'key':
                    command = f'input keyevent {command_data.get("key_code")}'
                else:
                    logger.warning(f"未知的命令类型: {command_type}")
                    results[device_id] = {'success': False, 'message': '未知的命令类型'}
                    continue
                
                # 执行命令
                result = self.device_manager.execute_command(device_id, command)
                results[device_id] = {
                    'success': result['success'],
                    'stdout': result.get('stdout'),
                    'stderr': result.get('stderr')
                }
            except Exception as e:
                logger.error(f"在设备 {device_id} 上执行命令时发生错误: {str(e)}")
                results[device_id] = {'success': False, 'message': str(e)}
        
        return results